package org.maitaole.interceptor;

import com.google.common.util.concurrent.RateLimiter;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class RateLimitInterceptor implements HandlerInterceptor {

    private Map<String, RateLimiter> ipToRateLimiter = new ConcurrentHashMap<>();

    private Integer rateLimit;

    public void setRateLimit(Integer rateLimit) {
        this.rateLimit = rateLimit;
    }

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
//        rateLimiter.tryAcquire();  // 尝试获取令牌
        String ipAddress = request.getRemoteHost();  // 获取远程的ip地址

        RateLimiter rl = ipToRateLimiter.get(ipAddress); // 获取对应的限流器
        if(null == rl) {
            rl = RateLimiter.create(rateLimit); // 每秒三个令牌
            ipToRateLimiter.put(ipAddress, rl);  // 重新设置
            rl.tryAcquire();
            return true;
        }else {
            if(rl.tryAcquire()) {
                return true;
            }
            // 429 表示用户的请求过快
            response.setStatus(429);
            return false;
        }
    }
}
